{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import os, json, time, ast, math, getpass\n",
        "from dataclasses import dataclass, field\n",
        "from typing import Dict, List, Callable, Any\n",
        "import google.generativeai as genai\n",
        "\n",
        "try:\n",
        "    import networkx as nx\n",
        "except ImportError:\n",
        "    nx = None"
      ],
      "metadata": {
        "id": "-X3WmICe5myt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def make_model(api_key: str, model_name: str = \"gemini-1.5-flash\"):\n",
        "    genai.configure(api_key=api_key)\n",
        "    return genai.GenerativeModel(model_name, system_instruction=(\n",
        "        \"You are GraphAgent, a principled planner-executor. \"\n",
        "        \"Prefer structured, concise outputs; use provided tools when asked.\"\n",
        "    ))\n",
        "\n",
        "def call_llm(model, prompt: str, temperature=0.2) -> str:\n",
        "    r = model.generate_content(prompt, generation_config={\"temperature\": temperature})\n",
        "    return (r.text or \"\").strip()"
      ],
      "metadata": {
        "id": "4PgAdOT05q13"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def safe_eval_math(expr: str) -> str:\n",
        "    node = ast.parse(expr, mode=\"eval\")\n",
        "    allowed = (ast.Expression, ast.BinOp, ast.UnaryOp, ast.Num, ast.Constant,\n",
        "               ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Pow, ast.Mod,\n",
        "               ast.USub, ast.UAdd, ast.FloorDiv, ast.AST)\n",
        "    def check(n):\n",
        "        if not isinstance(n, allowed): raise ValueError(\"Unsafe expression\")\n",
        "        for c in ast.iter_child_nodes(n): check(c)\n",
        "    check(node)\n",
        "    return str(eval(compile(node, \"<math>\", \"eval\"), {\"__builtins__\": {}}, {}))\n",
        "\n",
        "DOCS = [\n",
        "    \"Solar panels convert sunlight to electricity; capacity factor ~20%.\",\n",
        "    \"Wind turbines harvest kinetic energy; onshore capacity factor ~35%.\",\n",
        "    \"RAG = retrieval-augmented generation joins search with prompting.\",\n",
        "    \"LangGraph enables cyclic graphs of agents; good for tool orchestration.\",\n",
        "]\n",
        "def search_docs(q: str, k: int = 3) -> List[str]:\n",
        "    ql = q.lower()\n",
        "    scored = sorted(DOCS, key=lambda d: -sum(w in d.lower() for w in ql.split()))\n",
        "    return scored[:k]"
      ],
      "metadata": {
        "id": "0jrDj0nz5tDO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@dataclass\n",
        "class State:\n",
        "    task: str\n",
        "    plan: str = \"\"\n",
        "    scratch: List[str] = field(default_factory=list)\n",
        "    evidence: List[str] = field(default_factory=list)\n",
        "    result: str = \"\"\n",
        "    step: int = 0\n",
        "    done: bool = False\n",
        "\n",
        "def node_plan(state: State, model) -> str:\n",
        "    prompt = f\"\"\"Plan step-by-step to solve the user task.\n",
        "Task: {state.task}\n",
        "Return JSON: {{\"subtasks\": [\"...\"], \"tools\": {{\"search\": true/false, \"math\": true/false}}, \"success_criteria\": [\"...\"]}}\"\"\"\n",
        "    js = call_llm(model, prompt)\n",
        "    try:\n",
        "        plan = json.loads(js[js.find(\"{\"): js.rfind(\"}\")+1])\n",
        "    except Exception:\n",
        "        plan = {\"subtasks\": [\"Research\", \"Synthesize\"], \"tools\": {\"search\": True, \"math\": False}, \"success_criteria\": [\"clear answer\"]}\n",
        "    state.plan = json.dumps(plan, indent=2)\n",
        "    state.scratch.append(\"PLAN:\\n\"+state.plan)\n",
        "    return \"route\"\n",
        "\n",
        "def node_route(state: State, model) -> str:\n",
        "    prompt = f\"\"\"You are a router. Decide next node.\n",
        "Context scratch:\\n{chr(10).join(state.scratch[-5:])}\n",
        "If math needed -> 'math', if research needed -> 'research', if ready -> 'write'.\n",
        "Return one token from [research, math, write]. Task: {state.task}\"\"\"\n",
        "    choice = call_llm(model, prompt).lower()\n",
        "    if \"math\" in choice and any(ch.isdigit() for ch in state.task):\n",
        "        return \"math\"\n",
        "    if \"research\" in choice or not state.evidence:\n",
        "        return \"research\"\n",
        "    return \"write\"\n",
        "\n",
        "def node_research(state: State, model) -> str:\n",
        "    prompt = f\"\"\"Generate 3 focused search queries for:\n",
        "Task: {state.task}\n",
        "Return as JSON list of strings.\"\"\"\n",
        "    qjson = call_llm(model, prompt)\n",
        "    try:\n",
        "        queries = json.loads(qjson[qjson.find(\"[\"): qjson.rfind(\"]\")+1])[:3]\n",
        "    except Exception:\n",
        "        queries = [state.task, \"background \"+state.task, \"pros cons \"+state.task]\n",
        "    hits = []\n",
        "    for q in queries:\n",
        "        hits.extend(search_docs(q, k=2))\n",
        "    state.evidence.extend(list(dict.fromkeys(hits)))\n",
        "    state.scratch.append(\"EVIDENCE:\\n- \" + \"\\n- \".join(hits))\n",
        "    return \"route\"\n",
        "\n",
        "def node_math(state: State, model) -> str:\n",
        "    prompt = \"Extract a single arithmetic expression from this task:\\n\"+state.task\n",
        "    expr = call_llm(model, prompt)\n",
        "    expr = \"\".join(ch for ch in expr if ch in \"0123456789+-*/().%^ \")\n",
        "    try:\n",
        "        val = safe_eval_math(expr)\n",
        "        state.scratch.append(f\"MATH: {expr} = {val}\")\n",
        "    except Exception as e:\n",
        "        state.scratch.append(f\"MATH-ERROR: {expr} ({e})\")\n",
        "    return \"route\"\n",
        "\n",
        "def node_write(state: State, model) -> str:\n",
        "    prompt = f\"\"\"Write the final answer.\n",
        "Task: {state.task}\n",
        "Use the evidence and any math results below, cite inline like [1],[2].\n",
        "Evidence:\\n{chr(10).join(f'[{i+1}] '+e for i,e in enumerate(state.evidence))}\n",
        "Notes:\\n{chr(10).join(state.scratch[-5:])}\n",
        "Return a concise, structured answer.\"\"\"\n",
        "    draft = call_llm(model, prompt, temperature=0.3)\n",
        "    state.result = draft\n",
        "    state.scratch.append(\"DRAFT:\\n\"+draft)\n",
        "    return \"critic\"\n",
        "\n",
        "def node_critic(state: State, model) -> str:\n",
        "    prompt = f\"\"\"Critique and improve the answer for factuality, missing steps, and clarity.\n",
        "If fix needed, return improved answer. Else return 'OK'.\n",
        "Answer:\\n{state.result}\\nCriteria:\\n{state.plan}\"\"\"\n",
        "    crit = call_llm(model, prompt)\n",
        "    if crit.strip().upper() != \"OK\" and len(crit) > 30:\n",
        "        state.result = crit.strip()\n",
        "        state.scratch.append(\"REVISED\")\n",
        "    state.done = True\n",
        "    return \"end\"\n",
        "\n",
        "NODES: Dict[str, Callable[[State, Any], str]] = {\n",
        "    \"plan\": node_plan, \"route\": node_route, \"research\": node_research,\n",
        "    \"math\": node_math, \"write\": node_write, \"critic\": node_critic\n",
        "}\n",
        "\n",
        "def run_graph(task: str, api_key: str) -> State:\n",
        "    model = make_model(api_key)\n",
        "    state = State(task=task)\n",
        "    cur = \"plan\"\n",
        "    max_steps = 12\n",
        "    while not state.done and state.step < max_steps:\n",
        "        state.step += 1\n",
        "        nxt = NODES[cur](state, model)\n",
        "        if nxt == \"end\": break\n",
        "        cur = nxt\n",
        "    return state\n",
        "\n",
        "def ascii_graph():\n",
        "    return \"\"\"\n",
        "START -> plan -> route -> (research <-> route) & (math <-> route) -> write -> critic -> END\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "eBYrj_Eb5uI_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 882
        },
        "id": "_ID2dD6Q5HHn",
        "outputId": "51ea467e-b8c5-4aa7-8eef-ba9373a8c2db"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "🔐 Enter GEMINI_API_KEY: ··········\n",
            "📝 Enter your task: \n",
            "\n",
            "=== GRAPH === \n",
            "START -> plan -> route -> (research <-> route) & (math <-> route) -> write -> critic -> END\n",
            "\n",
            "\n",
            "✅ Result in 11.45s:\n",
            "The answer is factually incorrect in its claim of wind power being *more* reliable than solar power based solely on capacity factors. While it correctly states the approximate capacity factors, it omits crucial context.  Capacity factor is only one aspect of reliability;  intermittency and predictability are also vital. Solar power's output is more predictable during daylight hours, while wind power's output is more variable due to fluctuating wind speeds.  The inclusion of \"5 * 7 = 35\" is irrelevant and detracts from the answer.\n",
            "\n",
            "\n",
            "Improved Answer:\n",
            "\n",
            "Wind and solar power are both intermittent renewable energy sources, but differ in their reliability profiles.  While onshore wind power generally exhibits a higher capacity factor (around 35%) than solar photovoltaic (around 20%), this alone doesn't fully capture reliability.  Solar power's output is more predictable during daylight hours, making it easier to forecast and integrate into the grid. Wind power's output is more variable due to unpredictable wind speeds.  A comprehensive assessment of reliability requires considering factors beyond capacity factor, including geographic location, technological advancements (e.g., improved forecasting models), and grid integration strategies.  Therefore, a simple comparison based solely on capacity factor is insufficient to definitively declare one technology more reliable than the other.\n",
            "\n",
            "---- Evidence ----\n",
            "Solar panels convert sunlight to electricity; capacity factor ~20%.\n",
            "Wind turbines harvest kinetic energy; onshore capacity factor ~35%.\n",
            "\n",
            "---- Scratch (last 5) ----\n",
            "PLAN:\n",
            "{\n",
            "  \"subtasks\": [\n",
            "    \"Search for information comparing the reliability of solar and wind power generation.\",\n",
            "    \"Summarize the findings regarding reliability, considering factors like intermittency, predictability, and capacity factors.\",\n",
            "    \"Compute 5 * 7\"\n",
            "  ],\n",
            "  \"tools\": {\n",
            "    \"search\": true,\n",
            "    \"math\": true\n",
            "  },\n",
            "  \"success_criteria\": [\n",
            "    \"A summary comparing the reliability of solar and wind power is produced.\",\n",
            "    \"The calculation 5 * 7 is correctly performed and the result is reported.\"\n",
            "  ]\n",
            "}\n",
            "EVIDENCE:\n",
            "- Solar panels convert sunlight to electricity; capacity factor ~20%.\n",
            "- Wind turbines harvest kinetic energy; onshore capacity factor ~35%.\n",
            "- Wind turbines harvest kinetic energy; onshore capacity factor ~35%.\n",
            "- Solar panels convert sunlight to electricity; capacity factor ~20%.\n",
            "- Solar panels convert sunlight to electricity; capacity factor ~20%.\n",
            "- Wind turbines harvest kinetic energy; onshore capacity factor ~35%.\n",
            "DRAFT:\n",
            "Wind power is more reliable than solar power based on capacity factors [1, 2].  Onshore wind has a capacity factor of ~35%, while solar has a capacity factor of ~20% [1, 2].  5 * 7 = 35.\n",
            "REVISED\n"
          ]
        }
      ],
      "source": [
        "if __name__ == \"__main__\":\n",
        "    key = os.getenv(\"GEMINI_API_KEY\") or getpass.getpass(\"🔐 Enter GEMINI_API_KEY: \")\n",
        "    task = input(\"📝 Enter your task: \").strip() or \"Compare solar vs wind for reliability; compute 5*7.\"\n",
        "    t0 = time.time()\n",
        "    state = run_graph(task, key)\n",
        "    dt = time.time() - t0\n",
        "    print(\"\\n=== GRAPH ===\", ascii_graph())\n",
        "    print(f\"\\n✅ Result in {dt:.2f}s:\\n{state.result}\\n\")\n",
        "    print(\"---- Evidence ----\")\n",
        "    print(\"\\n\".join(state.evidence))\n",
        "    print(\"\\n---- Scratch (last 5) ----\")\n",
        "    print(\"\\n\".join(state.scratch[-5:]))"
      ]
    }
  ]
}